package com.ayuer.chapter06;

import java.util.ArrayList;

public class BinarySearchTree<Key extends Comparable<Key>, Value> {

    BSTNode<Key, Value> root;

    //根据key获取Value
    public Value get(Key key) {
        if (key == null) {
            return null;
        }
        BSTNode<Key, Value> curr = root;
        while (curr != null) {
            int compRes = key.compareTo(curr.key);
            if (key.compareTo(curr.key) == 0) {
                return curr.value;
            } else if (compRes < 0) {
                curr = curr.left;
            } else if (compRes > 0) {
                curr = curr.right;
            }
        }
        return null;
    }

    //查找最小key
    //往左走到头就是最小的key
    public Value min() {
        return min(root);
    }

    public Value min(BSTNode<Key, Value> node) {
        if (node == null) {
            return null;
        }
        BSTNode<Key, Value> curr = node;
        BSTNode<Key, Value> near = null;
        while (curr != null) {
            near = curr;
            curr = curr.left;
        }
        return near.value;
    }

    //查找最大key
    //往右走到头就是最大的key
    public Value max() {
        return max(root);
    }

    public Value max(BSTNode<Key, Value> node) {
        if (node == null) {
            return null;
        }
        BSTNode<Key, Value> curr = node;
        BSTNode<Key, Value> near = null;
        while (curr != null) {
            near = curr;
            curr = curr.right;
        }
        return near.value;
    }


    //添加新节点
    public void put(Key key, Value value) {
        if (key == null || value == null) {
            return;
        }
        if (root == null) {
            root = new BSTNode<>(key, value);
            return;
        }
        //1.key存在则更新
        BSTNode<Key, Value> curr = root;
        BSTNode<Key, Value> near = root;
        while (curr != null) {
            near = curr;
            int compRes = key.compareTo(curr.key);
            if (key.compareTo(curr.key) == 0) {
                curr.value = value;
                return;
            } else if (compRes < 0) {
                curr = curr.left;
            } else if (compRes > 0) {
                curr = curr.right;
            }
        }
        //2.key不存在就加入
        BSTNode<Key, Value> node = new BSTNode<>(key, value);
        if (near.key.compareTo(key) > 0) {
            near.left = node;
        } else {
            near.right = node;
        }
    }


    /**
     * 查找前驱节点
     * 1.节点有左子树，此时前驱节点就是左子树最大值
     * 2.节点没有左子树，则离他的最近祖先，自左而来的就是前驱节点
     *
     * @param key
     * @return
     */
    public Value precessor(Key key) {
        if (key == null) {
            return null;
        }
        //1.找到key节点
        BSTNode<Key, Value> curr = root;
        BSTNode<Key, Value> near = null;
        while (curr != null) {
            int compRes = key.compareTo(curr.key);
            if (key.compareTo(curr.key) == 0) {
                break;
            } else if (compRes < 0) {
                curr = curr.left;
            } else if (compRes > 0) {
                near = curr;
                curr = curr.right;
            }
        }

        //没找到该节点，返回null
        if (curr == null) {
            return null;
        }

        //2.1左子树不为空
        if (curr.left != null) {
            return max(curr.left);
        }

        //2.2左子树为空，则离他的最近祖先，自左而来的就是前驱节点
        return near != null ?
                near.value : null;
    }


    /**
     * 查找后继节点
     * 1.节点有右子树，此时后继节点就是右子树最小值
     * 2.节点没有右子树，则离它最近的祖先，自右边而来的就是后继节点
     *
     * @param key
     * @return
     */
    public Value successor(Key key) {
        if (key == null) {
            return null;
        }
        //1.找到key节点
        BSTNode<Key, Value> curr = root;
        BSTNode<Key, Value> near = null;
        while (curr != null) {
            int compRes = key.compareTo(curr.key);
            if (key.compareTo(curr.key) == 0) {
                break;
            } else if (compRes < 0) {
                near = curr;
                curr = curr.left;
            } else if (compRes > 0) {
                curr = curr.right;
            }
        }

        //没找到该节点，返回null
        if (curr == null) {
            return null;
        }

        //2.1右子树不为空
        if (curr.right != null) {
            return min(curr.right);
        }

        //2.右子树为空，则离他的最近祖先，自左而来的就是前驱节点
        return near != null ?
                near.value : null;
    }

    //删除节点
    public Value remove(Key key) {
        if (key == null) {
            return null;
        }
        ArrayList<Value> res = new ArrayList<>();
        //1.找到key节点
        BSTNode<Key, Value> curr = root;
        BSTNode<Key, Value> parent = null;
        while (curr != null) {
            int compRes = key.compareTo(curr.key);
            if (key.compareTo(curr.key) == 0) {
                break;
            } else if (compRes < 0) {
                parent = curr;
                curr = curr.left;
            } else if (compRes > 0) {
                parent = curr;
                curr = curr.right;
            }
        }

        //没找到该节点，返回null
        if (curr == null) {
            return null;
        }
        res.add(curr.value);
        //删除
        if (curr.left == null) {
            //1，3
            //被删除节点没有左孩子，有右孩子
            shift(parent, curr, curr.right);
        } else if (curr.right == null) {
            //2，3
            //被删除节点有左孩子，没有右孩子
            shift(parent, curr, curr.left);
        } else {
            //4
            //4.1找到被删除节点的后继节点
            //p就是后继节点
            BSTNode<Key, Value> p = curr.right;
            //后继节点的父节点
            BSTNode<Key, Value> postParent = null;
            while (p.left != null) {
                postParent = p;
                p = p.left;
            }
            //4.2被删除节点和后继节点的孩子节点
            //4.3处理后继的的孩子节点
            if (curr.right == p) {
                //相邻，curr一定是只剩下right
                shift(parent, curr, curr.right);
                p.left = curr.left;
            } else {
                //不相邻
                shift(postParent, curr, curr.right);
                shift(parent, curr, curr.right);
                p.right = curr.right;
                p.left = curr.left;
            }
        }

        return res.get(0);
    }

    //顶替移除
    private void shift(BSTNode<Key, Value> parent, BSTNode<Key, Value> deleted, BSTNode<Key, Value> child) {
        //删除的是根节点
        if (parent == null) {
            root = child;
            return;
        } else if (parent.left == deleted) {//被删除的节点是父节点的左孩子
            parent.left = child;
        } else if (parent.right == deleted) {//被删除的节点是父节点的右孩子
            parent.right = child;
        }
    }

    static class BSTNode<Key, Value> {
        Key key;
        Value value;
        BSTNode<Key, Value> left, right;

        public BSTNode(Value value) {
            this.value = value;
        }

        public BSTNode(Key key, Value value) {
            this.value = value;
            this.key = key;
        }

        public BSTNode(Key key, Value value, BSTNode<Key, Value> left, BSTNode<Key, Value> right) {
            this.value = value;
            this.key = key;
            this.left = left;
            this.right = right;
        }
    }

}
